import { OpenAIClient, AzureKeyCredential } from '@azure/openai'
import { type CreateMessage, OpenAIStream, StreamingTextResponse, ToolCallPayload, JSONValue } from 'ai'

import { tools } from '@/lib/chat-tools/chat-tools-data'
import { applyToolCall } from '@/lib/chat-tools/chat-tools'

//Reference: https://sdk.vercel.ai/docs/guides/providers/openai

const client = new OpenAIClient('https://dolphin-chat.openai.azure.com/', new AzureKeyCredential(process.env.AZURE_OPENAI_API_KEY!))

export const runtime = 'edge'
export async function POST(req: Request) {
    const model = 'gpt-35-turbo'
    const { messages } = await req.json()
    const response = await client.streamChatCompletions(model, messages, { tools, toolChoice: 'auto' })

    const stream = OpenAIStream(response, {
        experimental_onToolCall: async (call: ToolCallPayload, appendToolCallMessage: (result?: { tool_call_id: string; function_name: string; tool_call_result: JSONValue; }) => CreateMessage[]) => {
            const result = applyToolCall(call)
            if (result) {
                console.log('chat-with-tools.experimental_onToolCall -----> result: ', result)
                return client.streamChatCompletions(model, [...messages, ...appendToolCallMessage(result)], { tools, toolChoice: 'auto' })
            }
        },

        onStart: async () => {
            // This callback is called when the stream starts
            // You can use this to save the prompt to your database
            // await savePromptToDatabase(prompt);
        },

        onToken: async (token: string) => {
            // This callback is called for each token in the stream
            // You can use this to debug the stream or save the tokens to your database
            // console.log('token: ', token)
        },

        onFinal: (completion: string) => {
            // console.log('chat-with-tools.onFinal --------> completion: ', completion)
        },

        onCompletion: async (completion: string) => {
            // This callback is called when the stream completes
            // You can use this to save the final completion to your database
            // await saveCompletionToDatabase(completion);
        },

        // experimental_streamData: true
    })
    return new StreamingTextResponse(stream)
}